package failover

import (
	"context"
	"errors"
	"gitee.com/flycash/ai-gateway-demo/internal/service/llm"
	"sync/atomic"
	"time"
)

var ErrNotAvailableService = errors.New("没有可用节点")

type CurServiceProvider struct {
	service llm.Service
	cnt     *atomic.Int32
}

// 1. 标记一次调用失败，如果连续失败次数 > 3 就将服务标记为不可用
func (s *CurServiceProvider) MarkFailure() {
	s.cnt.Add(1)
}

// 2. 标记成功，将连续失败次数置为0，服务标记为可用
func (s *CurServiceProvider) MarkSuccess() {
	s.cnt.Store(0)
}

// 3. 获取当前服务的状态，返回可用不可用
func (s *CurServiceProvider) IsAvailable() bool {
	// 连续不可用大于 3
	return s.cnt.Load() > 3
}

func (s *CurServiceProvider) HealthCheck() bool {
	// 简单使用这个接口检查可不可以访问
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*3)
	defer cancel()
	return s.service.Healthy(ctx)
}

// ServiceGroup 只有初始化的时候才会增加节点
// 某个大模型对应的服务商
type ServiceGroup struct {
	services []*CurServiceProvider
	count    *atomic.Int64
}

func (s *ServiceGroup) GetAvailableServices() (*CurServiceProvider, error) {
	if len(s.services) == 0 {
		return nil, ErrNotAvailableService
	}
	count := s.count.Add(1)

	for i := 0; i < len(s.services); i++ {
		idx := (int(count) + i) % len(s.services)
		svc := s.services[idx]
		if svc.IsAvailable() {
			return svc, nil
		}
	}
	return nil, ErrNotAvailableService
}

func (s *ServiceGroup) HealthyCheck() {
	for _, svc := range s.services {
		if !svc.IsAvailable() {
			if svc.HealthCheck() {
				// 标记我这个节点又是可用了
				svc.MarkSuccess()
			}
		}
	}
}
